package twoD.IsoPfem.Plasticity;

import inf.jlinalg.ArrayVector;
import inf.jlinalg.BLAM;
import inf.jlinalg.GeneralMatrixLSESolver;
import inf.jlinalg.ILSESolver;
import inf.jlinalg.IMatrix;
import inf.jlinalg.MatrixFormat;
import inf.jlinalg.QuadraticMatrixInfo;
import inf.jlinalg.SolveFailedException;
import inf.text.ArrayFormat;

import java.util.ArrayList;

import oneD.fem.algorithm.MyLinalg;
import twoD.IsoPfem.IStructure;
import twoD.IsoPfem.Control.IConstraintCondition;
import twoD.IsoPfem.Control.IController;
import twoD.IsoPfem.Control.LoadingState;

public class ArcLengthController implements IController {

	private static final double TOL = 1e-4;

	private IConstraintCondition cc_;
	private int numStep_;
	private IStructure struct_;

	private ArrayList<LoadingState> state_;

	public ArcLengthController(IStructure struct, IConstraintCondition cc) {
		struct_ = struct;
		cc_ = cc;
		numStep_ = 1;
		state_ = new ArrayList<LoadingState>();
	}

	@Override
	public void Iterate() throws SolveFailedException {
		NR();
	}

	@Override
	public void setNumberOfStep(int n) {
		numStep_ = n;
	}

	@Override
	public int getNumberOfState() {
		return state_.size();
	}

	@Override
	public LoadingState getState(int i) {
		return state_.get(i);
	}

	@Override
	public IStructure getStructure() {
		return struct_;
	}

	public void listingState() {
		for (LoadingState s : state_) {
			System.out.println(s.getLambda());
			System.out.println(MatrixFormat.format(s.getU()));
		}
	}

	/**
	 * normal Newton-Raphson iteration scheme for arc length
	 * 
	 * @category reference from Prof. Meschke slides
	 * 
	 */
	private void NR() throws SolveFailedException {
		int NEQ = struct_.enumerateDOFs();
		// solver initialization
		ILSESolver solver = new GeneralMatrixLSESolver(
				GeneralMatrixLSESolver.SOLVER_IMPL_JAVA);
		// ILSESolver solver = new SpoolesLSESolver();
		// using this for solver implementation in Java, it will be slow
		QuadraticMatrixInfo aInfo = solver.getAInfo();
		aInfo.setSize(NEQ);

		IMatrix ktGlobal = solver.getA();
		double[] rGlobal = new double[NEQ];
		double[] riGlobal = new double[NEQ];
		double[] uGlobal = new double[NEQ];
		double lambda = cc_.getInitialLambda();
		struct_.solveStaticEquilibrium(lambda);
		uGlobal = struct_.getU();
		System.out.println("uGLobal = " + ArrayFormat.format(uGlobal));

		struct_.assembleExternalForceVector(rGlobal);
		double norm_rGlobal = MyLinalg.norm2(rGlobal);

		double[] forward_u = uGlobal.clone();
		double forward_lambda = 0;
		double[] uGlobal_trial = uGlobal.clone();
		double lambda_trial = lambda;

		for (int i = 0; i < numStep_; i++) {
			System.out.println("/***** Analysis step " + (i + 1) + " *****/");
			System.out.println("-->Control loop to find lambda & displacement");

			double[] uGlobal_n0 = uGlobal.clone();
			double lambda_n0 = lambda;

			cc_.storeStepSize();
			boolean converged = false;
			while (!converged) {
				/*
				 * predictor step
				 */
				solver.initialize();
				BLAM.zero(ktGlobal);
				struct_.setDisplacement(uGlobal);
				struct_.assembleTangentStiffnessMatrix(ktGlobal, true);
				// System.out.println("norm(ktGlobal)_F = "
				// + BLAM.norm(BLAM.NORM_FROBENIUS, ktGlobal));
				double[] delta_u_lambda = rGlobal.clone();
				solver.solve(delta_u_lambda);
				double s0 = Math.sqrt(1 + Math.pow(
						MyLinalg.norm2(delta_u_lambda), 2));

				double forward_criteria = MyLinalg.dot(delta_u_lambda,
						forward_u) + forward_lambda;
				if (forward_criteria < 0)
					s0 = -s0;

				MyLinalg.add(cc_.getS() / s0, delta_u_lambda, uGlobal);
				lambda += cc_.getS() / s0;

				uGlobal_trial = uGlobal.clone();
				lambda_trial = lambda;

				// System.out.println("s0 = " + s0 + ", lambda predictor = "
				// + lambda);

				/*
				 * corrector step
				 */
				double eps = 1;
				int loop_cnt = 0;
				while (eps > TOL) {
					solver.initialize();
					struct_.setDisplacement(uGlobal);
					BLAM.zero(ktGlobal);
					struct_.assembleTangentStiffnessMatrix(ktGlobal, true);
					// System.out.println("ktGlobal = "
					// + MatrixFormat.format(ktGlobal));
					// System.out.println("uGlobal_n0 = "
					// + ArrayFormat.format(uGlobal_n0));
					// System.out.println("lambda_n0 = " + lambda_n0);
					MyLinalg.zero(riGlobal);
					struct_.assembleInternalForceVector(riGlobal);
					// System.out.println("rGlobal = "
					// + ArrayFormat.format(rGlobal));
					// System.out.println("riGlobal = "
					// + ArrayFormat.format(riGlobal));
					double[] delta_u_r = rGlobal.clone();
					MyLinalg.scale(lambda, delta_u_r);
					MyLinalg.add(-1.0, riGlobal, delta_u_r);
					solver.solve(delta_u_r);
					delta_u_lambda = rGlobal.clone();
					solver.solve(delta_u_lambda);

					double f = cc_.f(uGlobal, uGlobal_n0, lambda, lambda_n0);
					double[] dfdu = cc_.dfdu(uGlobal, uGlobal_n0, lambda,
							lambda_n0);
					double dfdlambda = cc_.dfdlambda(uGlobal, uGlobal_n0,
							lambda, lambda_n0);

					double tmp1 = MyLinalg.dot(dfdu, delta_u_r);
					double tmp2 = MyLinalg.dot(dfdu, delta_u_lambda);
					double delta_lambda = -(f + tmp1) / (tmp2 + dfdlambda);

					double[] delta_u = delta_u_r.clone();
					MyLinalg.add(delta_lambda, delta_u_lambda, delta_u);
					MyLinalg.add(1.0, delta_u, uGlobal);
					System.out.println("lambda = " + lambda
							+ "\tdelta_lambda = " + delta_lambda);
					lambda += delta_lambda;

					// compute the relative error
					struct_.setDisplacement(uGlobal);
					MyLinalg.zero(riGlobal);
					struct_.assembleInternalForceVector(riGlobal);
					double[] tmp = riGlobal.clone();
					MyLinalg.add(-lambda, rGlobal, tmp);
					f = cc_.f(uGlobal, uGlobal_n0, lambda, lambda_n0);
					// eps = (MyLinalg.norm2(tmp) + Math.abs(f)) / norm_rGlobal;
					eps = Math.abs(delta_lambda);
					// System.out.println("uGlobal_n0 = "
					// + ArrayFormat.format(uGlobal_n0));
					// System.out.println("lambda_n0 = " + lambda_n0);
					// System.out.println("f = " + f);
					System.out.println("  '->global equilibrium = "
							+ MyLinalg.norm2(tmp) + "\tconstraint = "
							+ Math.abs(f));

					// adaptive technique to modify step size when iteration not
					// converge
					if (loop_cnt < 100)
						loop_cnt++;
					else {
						// reset the value of uGlobal and lambda
						MyLinalg.copy(uGlobal_n0, uGlobal);
						lambda = lambda_n0;
						// reduce step size
						cc_.reduceStepSize();
						System.out.println("step size reduced 1/2");
						// reset the step
						break;
					}

					// update ths structure internal state if need (i.e
					// plasticity)
					struct_.updateInternalState();
				}
				if (eps < TOL)
					converged = true;
				// System.out.println("eps = " + eps);
				// System.out.println("constrainted u = "
				// + ArrayFormat.format(cc_.getU(uGlobal)));
				System.out.println(" '->converged, nloop = " + loop_cnt);
			}

			struct_.updateGlobalState();
			// do another loop to ensure global equilibrium
			System.out.println("post checking global equilibrium state");
			double eps = 1.0;
			while (eps > TOL) {
				solver.initialize();
				BLAM.zero(ktGlobal);
				MyLinalg.zero(riGlobal);
				// struct_.setDisplacement(uGlobal);
				struct_.assembleTangentStiffnessMatrix(ktGlobal, true);
				struct_.assembleInternalForceVector(riGlobal);
				double[] delta_u_r = rGlobal.clone();
				MyLinalg.scale(lambda, delta_u_r);
				MyLinalg.add(-1.0, riGlobal, delta_u_r);
				solver.solve(delta_u_r);
				double[] delta_u_lambda = rGlobal.clone();
				solver.solve(delta_u_lambda);

				double f = cc_.f(uGlobal, uGlobal_n0, lambda, lambda_n0);
				double[] dfdu = cc_
						.dfdu(uGlobal, uGlobal_n0, lambda, lambda_n0);
				double dfdlambda = cc_.dfdlambda(uGlobal, uGlobal_n0, lambda,
						lambda_n0);

				double tmp1 = MyLinalg.dot(dfdu, delta_u_r);
				double tmp2 = MyLinalg.dot(dfdu, delta_u_lambda);
				double delta_lambda = -(f + tmp1) / (tmp2 + dfdlambda);

				double[] delta_u = delta_u_r.clone();
				MyLinalg.add(delta_lambda, delta_u_lambda, delta_u);
				MyLinalg.add(1.0, delta_u, uGlobal);
				System.out.println("lambda = " + lambda
						+ "\tdelta_lambda = " + delta_lambda);
				lambda += delta_lambda;

				struct_.setDisplacement(uGlobal);

				MyLinalg.zero(riGlobal);
				struct_.assembleInternalForceVector(riGlobal);
				double[] tmp = riGlobal.clone();
				MyLinalg.add(-lambda, rGlobal, tmp);
				f = cc_.f(uGlobal, uGlobal_n0, lambda, lambda_n0);
				eps = Math.abs(delta_lambda);
				System.out.println("  '->global equilibrium = "
						+ MyLinalg.norm2(tmp) + "\tconstraint = "
						+ Math.abs(f));
				struct_.updateInternalState();
			}
			state_.add(new LoadingState(lambda,
					new ArrayVector(uGlobal.clone())));
			System.out
					.println("-->End of control loop to find lambda & displacement");
			System.out.println("-->System summary:");
			System.out.println("system lambda = " + lambda);

			forward_u = uGlobal.clone();
			MyLinalg.add(-1, uGlobal_trial, forward_u);
			forward_lambda = lambda - lambda_trial;
			cc_.restoreStepSize();
			cc_.increment();

			// check global equilibrium
			double[] ritest = new double[NEQ];
			struct_.assembleInternalForceVector(ritest);
			MyLinalg.add(-lambda, rGlobal, ritest);
			System.out.println("Global equilibrium = " + MyLinalg.norm2(ritest)
					/ norm_rGlobal);

			// check local equilibrium

			System.out.println("/***** End of step " + (i + 1) + " *****/\n");
		}
	}

	/**
	 * modified Newton-Raphson iteration scheme with combination to line search.
	 * Apply sphere constraint with psi=0
	 * 
	 * @category Crisfield paper: AN ARC-LENGTH METHOD INCLUDING LINE SEARCHES
	 *           AND ACCELERATIONS
	 * 
	 */
	/*
	 * this need to be debugged
	 */
	// private void mNR_sp_psi0_ls() throws SolveFailedException {
	// int NEQ = struct_.enumerateDOFs();
	// // solver initialization
	// ILSESolver solver = new GeneralMatrixLSESolver(
	// GeneralMatrixLSESolver.SOLVER_IMPL_JAVA);
	// // ILSESolver solver = new SpoolesLSESolver();
	// // using this for solver implementation in Java, it will be slow
	// QuadraticMatrixInfo aInfo = solver.getAInfo();
	// aInfo.setSize(NEQ);
	//
	// IMatrix ktGlobal = solver.getA();
	// double[] rGlobal = new double[NEQ];
	// double[] riGlobal = new double[NEQ];
	// double[] uGlobal = new double[NEQ];
	// double lambda = 0;
	// double delta_l = cc_.getS();
	//
	// struct_.assembleExternalForceVector(rGlobal);
	//
	// for (int i = 0; i < numStep_; i++) {
	// System.out.println("/***** Analysis step " + (i + 1) + " *****/");
	// System.out.println("-->Control loop to find lambda & displacement");
	//
	// // double[] forward_u = uGlobal.clone();
	// // double forward_lambda = 0;
	// // double[] uGlobal_trial = null;
	// // double lambda_trial = 0;
	//
	// solver.initialize();
	// struct_.setDisplacement(uGlobal);
	// BLAM.zero(ktGlobal);
	// struct_.assembleTangentStiffnessMatrix(ktGlobal, true);
	//
	// double[] delta_u = new double[NEQ];
	// double[] delta_T = rGlobal.clone();
	// solver.solve(delta_T);
	// double c1 = MyLinalg.dot(delta_T, delta_T);
	// double d1 = MyLinalg.dot(delta_T, delta_u);
	//
	// boolean converged = false;
	// double eta_i = 1.0;
	// double eps = 1.0;
	// int cnt = 0;
	// while (!converged) {
	// MyLinalg.zero(riGlobal);
	// struct_.assembleInternalForceVector(riGlobal);
	// double[] delta_i_bar = riGlobal.clone();
	// MyLinalg.scale(-1.0, delta_i_bar);
	// solver.solve(delta_i_bar);
	//
	// // find lambda_i+1
	// if (cnt == 0) {
	// // predictor step
	// double delta_lambda_1 = 0;
	// double a = Math.signum(MyLinalg.dot(delta_T, rGlobal));
	// delta_lambda_1 = a * delta_l / Math.sqrt(c1);
	// delta_u = delta_T.clone();
	// // double forward_criteria = MyLinalg.dot(delta_u,
	// // forward_u)
	// // + forward_lambda;
	// // if (forward_criteria < 0)
	// // delta_lambda_1 = -delta_lambda_1;
	// MyLinalg.scale(delta_lambda_1, delta_u);
	// lambda += delta_lambda_1;
	// // uGlobal_trial = uGlobal.clone();
	// // lambda_trial = lambda;
	// System.out.println("lambda predictor = " + lambda);
	// } else {
	// // corrector
	// double a1 = eta_i * c1;
	// double d2 = MyLinalg.dot(delta_T, delta_i_bar);
	// double a2 = 2 * d1 + 2 * eta_i * d2;
	// double d3 = MyLinalg.dot(delta_i_bar, delta_i_bar);
	// double d4 = MyLinalg.dot(delta_i_bar, delta_u);
	// double a3 = eta_i * d3 + 2 * d4
	// + MyLinalg.dot(delta_u, delta_u) - delta_l
	// * delta_l;
	// double Delta = a2 * a2 - 4 * a1 * a3;
	// double x1 = (-a2 - Math.sqrt(Delta)) / (2 * a1);
	// double x2 = (-a2 + Math.sqrt(Delta)) / (2 * a1);
	// double cos1 = 1 + eta_i / (delta_l * delta_l)
	// * (d4 + x1 * d1);
	// double cos2 = 1 + eta_i / (delta_l * delta_l)
	// * (d4 + x2 * d1);
	// if (cos1 * cos2 < 0) {
	// if (cos1 > 0)
	// lambda = x1;
	// else
	// lambda = x2;
	// } else {
	// if (cos1 >= 0) {
	// if (a2 != 0) {
	// double dis1 = Math.abs(x1 + a3 / a2);
	// double dis2 = Math.abs(x2 + a3 / a2);
	// if (dis1 < dis2)
	// lambda = x1;
	// else
	// lambda = x2;
	// } else {
	// lambda = x1;
	// }
	// } else
	// throw new SolveFailedException(
	// "step error: there's no acute angle to advance the step");
	// }
	// delta_u = delta_i_bar.clone();
	// MyLinalg.add(lambda, delta_T, delta_u);
	// System.out.println("lambda = " + lambda + ", epsilon = "
	// + eps);
	// }
	// cnt++;
	// MyLinalg.add(eta_i, delta_u, uGlobal);
	// struct_.setDisplacement(uGlobal);
	// // eps = MyLinalg.norm2(delta_u);
	// double[] tmp = new double[NEQ];
	// struct_.assembleInternalForceVector(tmp);
	// MyLinalg.add(-lambda, rGlobal, tmp);
	// eps = MyLinalg.norm2(tmp);
	// if (eps < TOL)
	// converged = true;
	// // end of corrector
	// // forward_u = uGlobal.clone();
	// // MyLinalg.add(-1, uGlobal_trial, forward_u);
	// // forward_lambda = lambda - lambda_trial;
	// System.out.println("lambda converged = " + lambda);
	// struct_.updateInternalState();
	// }
	// state_.add(new LoadingState(lambda,
	// new ArrayVector(uGlobal.clone())));
	// System.out.println("-->System summary:");
	// System.out.println("system lambda = " + lambda);
	// System.out.println("constrainted u = "
	// + ArrayFormat.format(cc_.getU(uGlobal)));
	// System.out.println("/***** End of step " + (i + 1) + " *****/\n");
	// struct_.updateGlobalState();
	// }
	// }
}
